import os
import shutil
from patient_information import find_patient_files,load_patient_data,get_grade,get_murmur
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from python_speech_features import logfbank
from sklearn.model_selection import StratifiedKFold
from tqdm import tqdm
import wave
import librosa.display
import librosa
import soundfile
from spafe.features.gfcc import erb_spectrogram
from spafe.utils.vis import show_spectrogram
from spafe.utils.preprocessing import SlidingWindow
from data_split_kfold import cut_copy_files
from data_split_kfold import *




out_directory="data_5fold_new"
cut_data_directory="5fold_cut"
#创建五个文件夹存放训练集测试集特征标签
for i in range(5):
    kfold_out_directory = os.path.join(out_directory, str(i+1) + "_fold")
    label_directory = os.path.join(kfold_out_directory, "label")
    feat_directory = os.path.join(kfold_out_directory, "feat")
    if not os.path.exists(kfold_out_directory):
        os.makedirs(kfold_out_directory)

    if not os.path.exists(feat_directory):
        os.makedirs(feat_directory)

    if not os.path.exists(label_directory):
        os.makedirs(label_directory)

    #选择该折的测试集
    test_data_file = os.path.join(str(i+1) + "_fold")
    print(test_data_file)
    files = os.listdir(cut_data_directory)#cut数据文件夹
    train_feat = np.array([[[]]])
    train_label = np.array([])
    for f in sorted(files):
        #选四份作为训练
        if (f != test_data_file):

            feat = Log_GF(os.path.join(cut_data_directory,f))
            label,_,_ = get_label(os.path.join(cut_data_directory, f))
            if(len(train_feat) == 1):
                train_feat = feat
                train_label = label
            else:
                train_feat = np.concatenate((train_feat,feat),axis=0)
                train_label = np.concatenate((train_label,label),axis=0)

            print(os.path.join(cut_data_directory,f),"放在训练集")

        #其中一份作为测试集
        else:

            test_feat = Log_GF(os.path.join(cut_data_directory,f))
            np.save(feat_directory + r'/test_loggamma.npy', test_feat)

            label, location, id = get_label(os.path.join(cut_data_directory, f))

            np.save(label_directory + r'/test_label.npy', label)

            np.save(label_directory + r'/test_location.npy', location)

            np.save(label_directory + r'/test_id.npy', id)

            test_index = get_index(os.path.join(cut_data_directory,f))

            np.save(label_directory + r'/test_index.npy', test_index)

            print(os.path.join(cut_data_directory,f),"放在测试集")



    print("第",str(i+1),"折特征提取完毕")
    np.save(feat_directory + r'/train_loggamma.npy', train_feat)
    np.save(label_directory + r'/train_label.npy', train_label)



